# -*- coding:utf-8 -*-

import copy
import numpy as np
import open3d as o3d
import math
import json

# =========== 超参数 ===========
threshold = 0.03  # ICP距离阈值:2cm


#========== 工具函数 ===========
def demo_crop_geometry(pcd=None):
    print("手动几何裁剪演示")
    print(
        "1) 按两次“Y”以将几何体与Y轴的负方向对齐"
    )
    print("2) 按“K”锁定屏幕并切换到选择模式")
    print("3) 拖动以选择矩形,")
    print("   或者使用ctrl+左键单击进行多边形选择")
    print("4) 按“C”获取选定的几何图形并保存")
    print("5) 按“F”切换到自由视图模式")
    # 可视化几何体供用户交互
    o3d.visualization.draw_geometries_with_editing([pcd])


def draw_registration_result(source, target, transformation):
    """绘制配准结果"""
    source_temp = copy.deepcopy(source)
    target_temp = copy.deepcopy(target)
    source_temp.paint_uniform_color([1, 0.706, 0])
    target_temp.paint_uniform_color([0, 0.651, 0.929])
    source_temp.transform(transformation)
    o3d.visualization.draw_geometries([source_temp, target_temp])


def pick_points(pcd):
    print("")
    print(
        "1) 请使用至少选择三个对应关系 [shift + 左击]"
    )
    print("   按 [shift + 右击] 撤销拾取的点")
    print("2) 拾取点后，按“Q”关闭窗口")
    vis = o3d.visualization.VisualizerWithEditing()
    vis.create_window()
    vis.add_geometry(pcd)
    # 激活窗口。此函数将阻止当前线程，直到窗口关闭。
    vis.run()  # 等待用户拾取点
    picked = vis.get_picked_points()
    vis.destroy_window()
    print("")
    return picked

def demo_manual_registration(source=None, target=None):
    print("手动ICP演示")
    # 加载点云
    print("手动对齐前两点云的可视化")
    draw_registration_result(source, target, np.identity(4))

    # # 从两点云中拾取点并建立对应关系
    
    picked_id_target = pick_points(target)  
    picked_id_source = pick_points(source)
    assert (len(picked_id_source) >= 3 and len(picked_id_target) >= 3)
    assert (len(picked_id_source) == len(picked_id_target))
    corr = np.zeros((len(picked_id_source), 2))
    corr[:, 0] = picked_id_source
    corr[:, 1] = picked_id_target

    # 利用对应关系估计粗变换
    print("使用用户给定的对应关系计算粗糙变换")
    p2p = o3d.pipelines.registration.TransformationEstimationPointToPoint()
    trans_init = p2p.compute_transformation(source, target,
                                            o3d.utility.Vector2iVector(corr))

    # 用于改善的点对点ICP
    print("执行点对点ICP改善")
    reg_p2p = o3d.pipelines.registration.registration_icp(
        source, target, threshold, trans_init,
        o3d.pipelines.registration.TransformationEstimationPointToPoint(with_scaling=True))
    draw_registration_result(source, target, reg_p2p.transformation)

    # 提取变换矩阵参数
    trans = reg_p2p.transformation
    # translation = trans[:3, 3]
    # # 旋转矩阵转欧拉角（ZYX顺序：roll, pitch, yaw）

    # def rotation_matrix_to_euler_angles(R):
    #     sy = math.sqrt(R[0,0] * R[0,0] + R[1,0] * R[1,0])
    #     singular = sy < 1e-6
    #     if not singular:
    #         x = math.atan2(R[2,1], R[2,2])
    #         y = math.atan2(-R[2,0], sy)
    #         z = math.atan2(R[1,0], R[0,0])
    #     else:
    #         x = math.atan2(-R[1,2], R[1,1])
    #         y = math.atan2(-R[2,0], sy)
    #         z = 0
    #     # 转为角度
    #     return math.degrees(x), math.degrees(y), math.degrees(z)

    # pitch, yaw, roll = rotation_matrix_to_euler_angles(trans[:3, :3])
    result = {
        "transformation":trans.tolist(),
    }
    print(json.dumps(result, indent=4))
    print("")
    return result


if __name__ == "__main__":
    # 加载点云
    pcd = o3d.io.read_point_cloud("./test/pcd1.pcd")
    demo_crop_geometry(pcd)
    source = o3d.io.read_point_cloud("./test/pcd2.pcd")
    target = o3d.io.read_point_cloud("./test/pcd3.pcd")
    demo_manual_registration(source, target)